package org.snail.proxy.server.netty;

import com.google.common.collect.Maps;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.util.Attribute;
import io.netty.util.AttributeKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.snail.proxy.proto.DeleteClientRequest;
import org.snail.proxy.proto.MessageId;
import org.snail.proxy.proto.NetType;
import org.snail.proxy.proto.Packet;
import org.snail.proxy.proto.ProxyMessage;
import org.snail.proxy.proto.ProxyRequest;
import org.snail.proxy.server.InitRunner;
import org.snail.proxy.server.config.NettyConfig;
import org.snail.proxy.server.config.ServerConfig;
import org.snail.proxy.server.netty.http.HttpProxyServer;
import org.snail.proxy.server.netty.tcp.TcpProxyServer;

import java.util.Map;
import java.util.concurrent.CountDownLatch;

@ChannelHandler.Sharable
public class ProxyClientHandler extends ChannelInboundHandlerAdapter {
    private static final Logger log = LoggerFactory.getLogger(ProxyClientHandler.class);
    private static final AttributeKey<TcpSender> SENDER_KEY = AttributeKey.valueOf("senderKey");
    private static final AttributeKey<NettyServer> SERVER_KEY = AttributeKey.valueOf("serverKey");

    private final Map<String, Channel> proxyChannels = Maps.newConcurrentMap();
    private final ServerConfig serverConfig;

    public ProxyClientHandler(ServerConfig serverConfig) {
        this.serverConfig = serverConfig;
    }

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msgObj) throws Exception {
        Channel channel = ctx.channel();
        TcpSender sender = channel.attr(SENDER_KEY).get();
        if (sender != null) {
            Packet packet = (Packet) msgObj;
            Attribute<NettyServer> attribute = channel.attr(SERVER_KEY);
            switch (packet.getCmd()) {
                case MessageId.PROXY_REQUEST_VALUE :
                    ProxyRequest proxyRequest = ProxyRequest.parseFrom(packet.getBytes());
                    sender.setPort(proxyRequest.getProxyPort());

                    if (attribute.get() == null) {
                        NettyConfig config = serverConfig.getServers().get(proxyRequest.getProxyPort());
                        NettyServer server = createServer(proxyRequest.getType(), config, channel);
                        InitRunner.startServer(server, new CountDownLatch(2));
                        attribute.set(server);
                        log.info("create proxy: {}[{}].", proxyRequest.getProxyPort(), proxyRequest.getType());
                    }
                    break;
                case MessageId.PROXY_MESSAGE_VALUE :
                    ProxyMessage proxyMessage = ProxyMessage.parseFrom(packet.getBytes());
                    proxyChannels.get(proxyMessage.getId()).writeAndFlush(proxyMessage.getData().toByteArray());
                    attribute.get().onSendMessage(proxyMessage.getId());
                    break;
                case MessageId.DELETE_CLIENT_REQUEST_VALUE :
                    DeleteClientRequest deleteClientRequest = DeleteClientRequest.parseFrom(packet.getBytes());
                    Channel proxyChannel = proxyChannels.remove(deleteClientRequest.getId());
                    if (proxyChannel != null) {
                        proxyChannel.closeFuture();
                        log.info("delete client: {} count={}.", deleteClientRequest.getId(), proxyChannels.size());
                    }
                    break;
            }
            log.info("Receive proxy client message: {}[{}].", packet.getCmd(), packet.getBytes().length);
        } else {
            channel.close();
        }
        super.channelRead(ctx, msgObj);
    }

    private NettyServer createServer(NetType type, NettyConfig config, Channel channel) {
        switch (type) {
        case TCP:
            return new TcpProxyServer(config, channel, this);
        case HTTP:
            return new HttpProxyServer(config, channel, this);
        default:
            throw new UnsupportedOperationException("createServer un know type: " + type);
        }
    }

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception {
        Channel channel = ctx.channel();
        Attribute<TcpSender> channelAttr = channel.attr(SENDER_KEY);
        if (channelAttr.get() == null) {
            TcpSender sender = new TcpSender(channel);
            channelAttr.set(sender);
            log.info("Connect proxy client.");
        }
        super.channelActive(ctx);
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        Channel channel = ctx.channel();
        TcpSender sender = channel.attr(SENDER_KEY).get();
        if (sender != null) {
            NettyServer server = channel.attr(SERVER_KEY).get();
            if (server != null) {
                server.close();
            }
            log.info("Disconnect proxy client.");
        }
        super.channelInactive(ctx);
    }

    public void addProxyChannel(String id, Channel channel) {
        proxyChannels.put(id, channel);
    }

    public void removeProxyChannel(String id) {
        proxyChannels.remove(id);
    }
}
